#include <mpi.h>
#include <stdio.h>
#include <string.h>
#include <math.h>

int compareInt (const void *a, const void *b) {
	return *(int *)a - *(int *)b;
}

void *par_bsearch(const void *key, const void *base, size_t nmemb,
		       size_t size, int (*compar)(const void *, const void *),
			   MPI_Comm comm, int threads) {

	char *s = (char *)base;				// main buffer
	char *s_end = NULL;					// ptr to last element of main buffer
	char *new_s = NULL;					// recv buffer
	char *tmp = NULL;					// temp for checking return vals
	int *globalIdx;						// global index of found value
	unsigned int i;						// iterator
	int status = 0;						// utility status variable
	int split;							// # of elements in this split
	int newSplit;						// # of elements in next split
	int src;							// source of next MPI send
	int rank, nprocs;					// this proc's ID and num procs
	int globalSplit;					// split for all procs except last
	int surplus;						// surplus assigned to last proc

	  // Get rank and size/numprocs
	MPI_Comm_size(comm ,&nprocs); 
	MPI_Comm_rank(comm ,&rank);
	
	unsigned int sendcnts[nprocs];		// # bytes to send on scatter
	unsigned int displs[nprocs];		// displacements for scatter

	  // Compute global and local split
	surplus = nmemb % nprocs;
	globalSplit = split = (int)floor(1.0 * nmemb / nprocs);
	if (rank == nprocs - 1) split += surplus;

	globalIdx = (int *)malloc(sizeof(int));
	if (globalIdx == NULL) {
		printf("Malloc error in %d\n", rank);
		MPI_Abort(comm, 0);
	}

	  // Check to see if this proc might have the key
	s_end = s + size * (split - 1);
	if (compar((const void *)key, (const void *)s) >= 0 &&
		compar((const void *)key, (const void *)(s_end)) <= 0) 
		status = rank + 1; // add 1 because procs are 0-indexed
	
	  // Reduce to get ID of candidate proc
	MPI_Allreduce(&status, &src, 1, MPI_INT, MPI_SUM, comm);

	  // No candidate, so exit. Decrement to bring back to 0-indexed
	if (src == 0) return NULL; src--; 

	  // Increment global index based on candidate proc
	*globalIdx = globalSplit * src;

	  // Make sure we have enough elements to make an MPI iteration worthwhile
	while (globalSplit > threads && globalSplit > nprocs) {

	  	  // If we're the candidate, prepare scatter parameters
		if (rank == src) {
			newSplit = (int)floor(split / nprocs);
			for (i = 0; i < nprocs; i++) {
				sendcnts[i] = newSplit * size;
				displs[i] = i * newSplit * size;
			}
			sendcnts[nprocs-1] += (split % nprocs) * size;

#ifdef DEBUG
			printf("Src = %d.  Rank = %d.  Split = %d\n", src, rank, split);
			for (i = 0; i < nprocs; i++) printf("\t%u = %u\n", i, sendcnts[i]);
#endif
		}

#ifdef DEBUG
		printf("rank = %d, key = %d, src = %d, first = %d, last = %d"
				" split = %d, newsplit = %d\n", 		
			rank, *(int *)key, src, *(int *)s, *(int *)(s_end),
			split, newSplit);
#endif

		  // Calculate new splits
		if (src == nprocs - 1) globalSplit += surplus;
		surplus = globalSplit % nprocs;
		globalSplit = split = (int)floor(globalSplit / nprocs);
		if (rank == nprocs - 1) split += surplus;

		  // Prepare recv buffer to get newSplit elements
		tmp = (char *)malloc(split * size);
		if (tmp == NULL) {
			printf("Malloc error in %d\n", rank);
			MPI_Abort(comm, 0);
		}
		new_s = tmp;

		  // Scatter portions of array
		MPI_Scatterv((void *)s, (int *)sendcnts, (int *)displs, MPI_CHAR, 
			(void *)new_s, split * size, MPI_CHAR, src, comm);

		if (s != base) free(s);	  	// done with old array, free it ... 
		s = new_s;					// and assign new array
		s_end = s + (split - 1) * size;

#ifdef DEBUG
		printf("rank = %d, key = %d, src = %d, first = %d, last = %d"
				" split = %d, newsplit = %d\n", 		
			rank, *(int *)key, src, *(int *)s, *(int *)(s_end),
			split, newSplit);
#endif
	  	  // Check to see if this proc might have the key
		status = 0;  // clean status
		if (compar((const void *)key, (const void *)s) >= 0 &&
			compar((const void *)key, (const void *)(s_end)) <= 0) 
			status = rank + 1;

	  	  // Reduce to get ID of candidate proc
		MPI_Allreduce(&status, &src, 1, MPI_INT, MPI_SUM, comm);

	      // No candidate, so exit.
		if (src == 0) return NULL; src--;

		  // Update global index	
		*globalIdx += globalSplit * src;
	}

	  // Out of MPI
	  //   Src continues to spawn OpenMP threads
	  //   All other procs wait for bcast of location
	if (rank == src) {
		status = -1; // clean status as -1 because 0 could be valid
#ifdef DEBUG
		printf("In OMP as %d with %d - %d\n", rank,
				*(int *)s, *(int *)(s + (split-1) * size));
#endif
		  // OMP parallel: can't iterate over pointers, return, or break
#pragma omp parallel for num_threads(threads)
		for (i = 0; i < split; i += 1) 
			if (compar((const void *)key, (const void *)(s + i * size)) == 0) 
				status = i;  // found key, save location

		  // Increment global index
		*globalIdx += status;

		  // return key location if found
		if (status != -1) status = *globalIdx;
	} else
		status = 0;

	  // Send result to all procs
	MPI_Allreduce(&status, globalIdx, 1, MPI_INT, MPI_SUM, comm);

	  // If we reduce to -1, we didn't find the value
	if (*globalIdx == -1) globalIdx = NULL;

	  // If safe, free up that memory
	if (s != NULL && s != base) free(s);

	return globalIdx;
}


int main(int argc, char *argv[])
{
    int key = 500002;				// value to find
	int n = 1000000;				// size of the array
	int i;							// iterator
	int rank, size;					// rank and size for MPI
	int split;						// size of each array portion
	int base;						// root of this procs array
	int *values = NULL;				// the array itself
	int *pItem = NULL;				// return value from par_bsearch

	if (argc > 1) key = atoi(argv[1]); // optional key on cmd line

	  // Init the MPI environment
	MPI_Init(&argc,&argv); 
	MPI_Comm comm = MPI_COMM_WORLD;
	MPI_Comm_size(comm ,&size); 
	MPI_Comm_rank(comm ,&rank);

	/*****************************************************
	 * Create a portion of the array on each proc
	 *****************************************************/
	split = (int)floor(1.0 * n / size);
	base = split * rank * 2;
	if (rank == size - 1)
		split += n % size;

	values = (int *)malloc(split * sizeof(int));
	for (i = 0; i < split; i++)	values[i] = (base) + i * 2;
	/*****************************************************/

	pItem = (int *)par_bsearch(&key, values, n, sizeof(int), 
		compareInt, comm, 16);

#ifdef DEBUG
	printf("%d is done\n", rank);
#endif

	if (rank == 0) 
		if (pItem != NULL) 
			printf ("%d is in the array at %d.\n", key, *pItem);
		else
			printf ("%d is not in the array.\n", key);

	MPI_Finalize();

    return 0;
}

